//
// Description: 863. 二叉树中所有距离为 K 的结点
// Created by Loading on 2021/7/28.
//

#include <bits/stdc++.h>
#include "Data_Structure_Definition.h"

using namespace std;

vector<int> distanceK(TreeNode *root, TreeNode *target, int k) {
    vector<int> res;
    // 存放每个节点值的父节点 （每个节点的值唯一，使用其值作为 key）
    unordered_map<int, TreeNode *> father;

    function<void(TreeNode *)> setFather = [&](TreeNode *node) {
        if (node->left) {
            father[node->left->val] = node;
            setFather(node->left);
        }
        if (node->right) {
            father[node->right->val] = node;
            setFather(node->right);
        }
    };

    // pre 为前驱节点
    function<void(TreeNode *, TreeNode *, int)> dfs = [&](TreeNode *node, TreeNode *pre, int depth) {
        if (!node) {
            return;
        }
        if (depth == k) {
            res.emplace_back(node->val);
            return;
        }
        /* 左子树、右子树与父节点，不为 node 的前驱节点时，递归寻找，避免重复计算 */
        if (node->left != pre) {
            dfs(node->left, node, depth + 1);
        }
        if (node->right != pre) {
            dfs(node->right, node, depth + 1);
        }
        if (father[node->val] != pre) {
            dfs(father[node->val], node, depth + 1);
        }
    };

    // 设置每个节点的父节点，根节点的父节点为 nullptr
    setFather(root);
    // DFS 寻找距离目标节点为 k 的节点
    dfs(target, nullptr, 0);

    return res;
}

int main() {
    auto *root = new TreeNode(3);

    root->left = new TreeNode(5);
    root->right = new TreeNode(1);

    root->left->left = new TreeNode(6);
    root->left->right = new TreeNode(2);
    root->right->left = new TreeNode(0);
    root->right->right = new TreeNode(8);

    root->left->right->left = new TreeNode(7);
    root->left->right->right = new TreeNode(4);

    int k = 2;
    vector<int> res = distanceK(root, root->left, k);

    for (auto &x : res) {
        cout << x << " ";
    }
    cout << endl;
}